#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <errno.h>

#include "include.h"
#include "HelperUtils.h"

#include "ftpclient.h"

void handleBye (int serverCommandSocket, const char* command, int* isExit);
void handleListMusic(int serverCommandSocket, const char* command);
void handleListPictures (int serverCommandSocket, const char* command);
void handlePut (int serverCommandSocket, const char* getCommand);
void handleGet (int serverCommandSocket, const char* getCommand);

int main(int argc, char** argv) {
    // Test for correct number of arguments
    if (argc != 3) {
        fprintf(stderr, "Usage: %s <Server Address> <Server Port>\n", argv[0]);
        exit(1);
    }
    
    // Process command line arguments
    char *serverIP = argv[1];
    int serverPort = atoi(argv[2]);

    // Create socket and connect to server
    int serverCommandSocket;
    if ((serverCommandSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0) {
        dieWithError("FATAL: socket() failed.");
    }
    
    struct sockaddr_in serverAddress;
    memset(&serverAddress, 0, sizeof(serverAddress));
    serverAddress.sin_family = AF_INET;
    serverAddress.sin_addr.s_addr = inet_addr(serverIP);
    serverAddress.sin_port = htons(serverPort);
    
    if (connect(serverCommandSocket, (struct sockaddr *) &serverAddress, sizeof(serverAddress)) < 0) {
        dieWithError("FATAL: connect() failed.");
    }

    // Get client command and send to server
    char *command = NULL;
    size_t commandSize = 0;
    int isExit = 0;
    while (!isExit) {
        command = NULL;
        commandSize = 0;
        fprintf(stdout, "> ");
        getLongLine(&command, &commandSize);
        if (strcmp(command, "BYE") == 0) {
            #if DEBUG == 1
            fprintf(stderr, "BYE\n");
            #endif
            handleBye(serverCommandSocket, command, &isExit);
        } else if (strcmp(command, "LIST-MUSIC") == 0) {
            #if DEBUG == 1
            fprintf(stderr, "LIST-MUSIC\n");
            #endif
            handleListMusic(serverCommandSocket, command);
        } else if (strcmp(command, "LIST-PICTURES") == 0) {
            #if DEBUG == 1
            fprintf(stderr, "LIST-PICTURES\n");
            #endif
            handleListPictures(serverCommandSocket, command);
        } else if (strncmp(command, "GET", 3) == 0) {
            #if DEBUG == 1
            fprintf(stderr, "GET\n");
            #endif
            handleGet(serverCommandSocket, command);
        } else if (strncmp(command, "PUT", 3) == 0) {
            #if DEBUG == 1
            fprintf(stderr, "PUT\n");
            #endif
            handlePut(serverCommandSocket, command);
        } else {
            fprintf(stdout, "Error: Unrecognized command.\n");
        }
        free(command);
    }

    return 0;
}

void handleBye (int serverCommandSocket, const char* command, int* isExit) {
    const char eot = 0x4;

    // Send bye to server
    if (send(serverCommandSocket, command, strlen(command), 0) != strlen(command))
        dieWithError("FATAL: send() failed.");
    // Send command terminator
    if (send(serverCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
        dieWithError("FATAL: send() failed.");

    // Close connection and exit
    close(serverCommandSocket);
    *isExit = 1;
}

void handleListMusic(int serverCommandSocket, const char* command) {
    char buffer[CMD_RES_SIZE];
    const char eot = 0x4;

    // Send list music to server
    if (send(serverCommandSocket, command, strlen(command), 0) != strlen(command))
        dieWithError("FATAL: send() failed.");
    // Send command terminator
    if (send(serverCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
        dieWithError("FATAL: send() failed.");

    // Receive response from server and display to stdout
    int byteReceived,
        eotReceived;
    byteReceived = recv(serverCommandSocket, buffer, CMD_RES_SIZE-1, 0);
    eotReceived = (buffer[byteReceived-1] == 0x4) ? 1 : 0;
    // Kill the EOT character
    if (eotReceived) buffer[byteReceived-1] = '\0';
    // NULL pad the string
    buffer[byteReceived] = '\0';
    // Print to stdout
    fprintf(stdout, "%s", buffer);
    // While last byte is not EOT
    while (!eotReceived) {
        byteReceived = recv(serverCommandSocket, buffer, CMD_RES_SIZE-1, 0);
        eotReceived = (buffer[byteReceived-1] == 0x4) ? 1 : 0;
        // Kill the EOT character
        if (eotReceived) buffer[byteReceived-1] = '\0';
        // NULL pad the string
        buffer[byteReceived] = '\0';
        // Print to stdout
        fprintf(stdout, "%s", buffer);
    }
}

void handleListPictures (int serverCommandSocket, const char* command) {
    char buffer[CMD_RES_SIZE];
    const char eot = 0x4;

    // Send list music to server
    if (send(serverCommandSocket, command, strlen(command), 0) != strlen(command))
        dieWithError("FATAL: send() failed.");
    // Send command terminator
    if (send(serverCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
        dieWithError("FATAL: send() failed.");

    // Receive response from server and display to stdout
    int byteReceived,
        eotReceived;
    byteReceived = recv(serverCommandSocket, buffer, CMD_RES_SIZE-1, 0);
    eotReceived = (buffer[byteReceived-1] == 0x4) ? 1 : 0;
    // Kill the EOT character
    if (eotReceived) buffer[byteReceived-1] = '\0';
    // NULL pad the string
    buffer[byteReceived] = '\0';
    // Print to stdout
    fprintf(stdout, "%s", buffer);
    // While last byte is not EOT
    while (!eotReceived) {
        byteReceived = recv(serverCommandSocket, buffer, CMD_RES_SIZE-1, 0);
        eotReceived = (buffer[byteReceived-1] == 0x4) ? 1 : 0;
        // Kill the EOT character
        if (eotReceived) buffer[byteReceived-1] = '\0';
        // NULL pad the string
        buffer[byteReceived] = '\0';
        // Print to stdout
        fprintf(stdout, "%s", buffer);
    }
}

void handlePut (int serverCommandSocket, const char* getCommand) {
    const char eot = 0x4;
    char sendDataBuffer[FILE_SEND_BUF_SIZE];

    // find out if PUT command is a valid one
    // format of PUT command is:
    //  PUT [MUSIC|PICTURES] [File Name]
    char *token1 = strchr(getCommand, ' ');
    char *token2;
    // token1 is NULL means no space is found in getCommand
    // token1 is less than 2 characters means subsequent part is malformatted
    // Let user know the command is malformatted
    if (token1 == NULL || strlen(token1) < 2) {
        fprintf(stdout, "Error: Malformatted PUT command.\n");
        return;
    } else {
        token2 = strchr(token1+1, ' ');
        if (token2 == NULL || strlen(token2) < 2) {
            fprintf(stdout, "Error: Malformatted PUT command.\n");
            return;
        }
    }

    // Extract file name and server path from tokens
    char *fileName = token2 + 1;
    char serverPath[16];
    memset(&serverPath, 0, sizeof(serverPath));
    strncpy(serverPath, token1+1, strlen(token1)-strlen(token2)-1);
    // If serverPath is not MUSIC or PICTURES, let user know
    if (strcmp(serverPath, "MUSIC")!=0 && strcmp(serverPath, "PICTURES")!=0) {
        fprintf(stdout, "Error: Unrecognized destination path \"%s\"\n", serverPath);
        return;
    }
    
    // Open file for read
    FILE* file = fopen(fileName, "r");
    if (file == NULL) {
        fprintf(stdout, "Error: Failed to open %s.\n", fileName);
        return;
    }

    // Listen on port DATA_PORT
    // socket(), bind(), listen()
    int clientDataSocket;
    int clientDataPort = DATA_PORT;
    struct sockaddr_in clientAddress;
    if ((clientDataSocket = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0)
        dieWithError("FATAL: socket() failed.");
    memset(&clientAddress, 0, sizeof(clientAddress));
    clientAddress.sin_family = AF_INET;
    clientAddress.sin_addr.s_addr = htonl(INADDR_ANY);
    clientAddress.sin_port = htons(clientDataPort);
    
    if (bind(clientDataSocket, (struct sockaddr *)&clientAddress, sizeof(clientAddress)) < 0)
        dieWithError("FATAL: bind() failed.");
    if (listen(clientDataSocket, MAXPENDING) < 0)
        dieWithError("FATAL: listen() failed.");

    // Let server know we are ready to send the file
    // Send command buffer length is length of getCommand + portNumber
    char* sendCommandBuffer = malloc(sizeof(getCommand)+32);
    // TODO: Check malloc return
    // Obtain data port in string
    char clientDataPortString[32];
    sprintf(clientDataPortString, "%d", clientDataPort);

    // Assemble the full command to be sent
    strcpy(sendCommandBuffer, "PUT ");
    strcat(sendCommandBuffer, clientDataPortString);
    strcat(sendCommandBuffer, " ");
    strcat(sendCommandBuffer, serverPath);
    strcat(sendCommandBuffer, "/");
    strcat(sendCommandBuffer, fileName);
    // Send the command to server
    if (send(serverCommandSocket, sendCommandBuffer, strlen(sendCommandBuffer), 0) != strlen(sendCommandBuffer))
        dieWithError("FATAL: send() failed.");
    // Append an EOT character to indicate end of transmission
    if (send(serverCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
        dieWithError("FATAL: send() failed.");

    free(sendCommandBuffer);

    // Check message returned by server
    char buffer[CMD_RES_SIZE];
    memset(&buffer, 0, sizeof(buffer));
    // Receive response from server and display to stdout
    int byteReceived,
        eotReceived,
        isError = 1;    // If server indicates an error
    byteReceived = recv(serverCommandSocket, buffer, CMD_RES_SIZE-1, 0);
    eotReceived = (buffer[byteReceived-1] == 0x4) ? 1 : 0;
    // Kill the EOT character
    if (eotReceived) {
        // If only byte received is EOT, server is telling us to send the file
        if (buffer[0] == 0x4)   isError = 0;
        // NULL pad the EOT
        buffer[byteReceived-1] = '\0';
    }
    // NULL pad the string
    buffer[byteReceived] = '\0';
    // Print to stdout
    fprintf(stdout, "%s", buffer);
    // While last byte is not EOT
    while (!eotReceived) {
        byteReceived = recv(serverCommandSocket, buffer, CMD_RES_SIZE-1, 0);
        eotReceived = (buffer[byteReceived-1] == 0x4) ? 1 : 0;
        // Kill the EOT character
        if (eotReceived) buffer[byteReceived-1] = '\0';
        // NULL pad the string
        buffer[byteReceived] = '\0';
        // Print to stdout
        fprintf(stdout, "%s", buffer);
    }
    
    if (!isError) {
        // Everthing looks good, Accept connection on data socket
        int serverDataSocket;
        struct sockaddr_in serverDataAddress;
        socklen_t serverDataAddressLength = sizeof(serverDataAddress);
        if ((serverDataSocket = accept(clientDataSocket, (struct sockaddr *)&serverDataAddress, &serverDataAddressLength)) < 0)
            dieWithError("FATAL: accept() failed.");

        // Send file in FILE_SEND_BUF_SIZE chunk
        int numberRead = 0;
        while ((numberRead = fread(sendDataBuffer, 1, FILE_SEND_BUF_SIZE, file)) == FILE_SEND_BUF_SIZE) {
            if (send(serverDataSocket, sendDataBuffer, numberRead, 0) != numberRead) 
                dieWithError("FATAL: send() failed.");
        }
        if (send(serverDataSocket, sendDataBuffer, numberRead, 0) != numberRead) 
            dieWithError("FATAL: send() failed.");

        close(serverDataSocket);
    }

    // Clean up
    fclose(file);
    close(clientDataSocket);
}

void handleGet (int serverCommandSocket, const char* getCommand) {
    const char eot = 0x4;
    char receiveDataBuffer[FILE_SEND_BUF_SIZE];

    // find out if GET command is a valid one
    // format of GET command is:
    //  GET [MUSIC|PICTURES] [File Name]
    char *token1 = strchr(getCommand, ' ');
    char *token2;
    // token1 is NULL means no space is found in getCommand
    // token1 is less than 2 characters means subsequent part is malformatted
    // Let user know the command is malformatted
    if (token1 == NULL || strlen(token1) < 2) {
        fprintf(stdout, "Error: Malformatted GET command.\n");
        return;
    } else {
        token2 = strchr(token1+1, ' ');
        if (token2 == NULL || strlen(token2) < 2) {
            fprintf(stdout, "Error: Malformatted GET command.\n");
            return;
        }
    }

    // Extract file name and server path from tokens
    char *fileName = token2 + 1;
    char serverPath[16];
    memset(&serverPath, 0, sizeof(serverPath));
    strncpy(serverPath, token1+1, strlen(token1)-strlen(token2)-1);
    // If serverPath is not MUSIC or PICTURES, let user know
    if (strcmp(serverPath, "MUSIC")!=0 && strcmp(serverPath, "PICTURES")!=0) {
        fprintf(stdout, "Error: Unrecognized destination path \"%s\"\n", serverPath);
        return;
    }
    
    // Open file for read
    FILE* file = fopen(fileName, "w");
    if (file == NULL) {
        fprintf(stdout, "Error: Failed to open %s.\n", fileName);
        return;
    }

    // Listen on port DATA_PORT
    // socket(), bind(), listen()
    int clientDataSocket;
    int clientDataPort = DATA_PORT;
    struct sockaddr_in clientAddress;
    if ((clientDataSocket = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0)
        dieWithError("FATAL: socket() failed.");
    memset(&clientAddress, 0, sizeof(clientAddress));
    clientAddress.sin_family = AF_INET;
    clientAddress.sin_addr.s_addr = htonl(INADDR_ANY);
    clientAddress.sin_port = htons(clientDataPort);
    
    if (bind(clientDataSocket, (struct sockaddr *)&clientAddress, sizeof(clientAddress)) < 0)
        dieWithError("FATAL: bind() failed.");
    if (listen(clientDataSocket, MAXPENDING) < 0)
        dieWithError("FATAL: listen() failed.");

    // Let server know we are ready to send the file
    // Send command buffer length is length of getCommand + portNumber
    char* sendCommandBuffer = malloc(sizeof(getCommand)+32);
    // TODO: Check malloc return
    // Obtain data port in string
    char clientDataPortString[32];
    sprintf(clientDataPortString, "%d", clientDataPort);

    // Assemble the full command to be sent
    strcpy(sendCommandBuffer, "GET ");
    strcat(sendCommandBuffer, clientDataPortString);
    strcat(sendCommandBuffer, " ");
    strcat(sendCommandBuffer, serverPath);
    strcat(sendCommandBuffer, "/");
    strcat(sendCommandBuffer, fileName);
    // Send the command to server
    if (send(serverCommandSocket, sendCommandBuffer, strlen(sendCommandBuffer), 0) != strlen(sendCommandBuffer))
        dieWithError("FATAL: send() failed.");
    // Append an EOT character to indicate end of transmission
    if (send(serverCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
        dieWithError("FATAL: send() failed.");

    free(sendCommandBuffer);

    // Check message returned by server
    char buffer[CMD_RES_SIZE];
    memset(&buffer, 0, sizeof(buffer));
    // Receive response from server and display to stdout
    int byteReceived,
        eotReceived,
        isError = 1;    // If server indicates an error
    byteReceived = recv(serverCommandSocket, buffer, CMD_RES_SIZE-1, 0);
    eotReceived = (buffer[byteReceived-1] == 0x4) ? 1 : 0;
    // Kill the EOT character
    if (eotReceived) {
        // If only byte received is EOT, server is telling us to send the file
        if (buffer[0] == 0x4)   isError = 0;
        // NULL pad the EOT
        buffer[byteReceived-1] = '\0';
    }
    // NULL pad the string
    buffer[byteReceived] = '\0';
    // Print to stdout
    fprintf(stdout, "%s", buffer);
    // While last byte is not EOT
    while (!eotReceived) {
        byteReceived = recv(serverCommandSocket, buffer, CMD_RES_SIZE-1, 0);
        eotReceived = (buffer[byteReceived-1] == 0x4) ? 1 : 0;
        // Kill the EOT character
        if (eotReceived) buffer[byteReceived-1] = '\0';
        // NULL pad the string
        buffer[byteReceived] = '\0';
        // Print to stdout
        fprintf(stdout, "%s", buffer);
    }
    
    if (!isError) {
        // Everthing looks good, Accept connection on data socket
        int serverDataSocket;
        struct sockaddr_in serverDataAddress;
        socklen_t serverDataAddressLength = sizeof(serverDataAddress);
        if ((serverDataSocket = accept(clientDataSocket, (struct sockaddr *)&serverDataAddress, &serverDataAddressLength)) < 0)
            dieWithError("FATAL: accept() failed.");

        // Receive file in FILE_SEND_BUF_SIZE chunk
        // After sending the whole file, server is expected to close connection.
        int byteReceived;
        int numberWritten = 0;
        while ((byteReceived = recv(serverDataSocket, receiveDataBuffer, FILE_SEND_BUF_SIZE, 0)) > 0) {
            // Write the buffer to file
            if ((numberWritten = fwrite(receiveDataBuffer, 1, byteReceived, file)) != byteReceived)
                dieWithError("FATAL: fwrite() failed.");
        }
        close(serverDataSocket);
    }

    // Clean up
    fclose(file);
    close(clientDataSocket);
}

